#!/bin/bash
set -e

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# Test images
TEST_IMAGES=(
    "ghcr.io/stargz-containers/nginx:1.25-estargz"
    "ghcr.io/stargz-containers/alpine:3.18-estargz"
    "ghcr.io/stargz-containers/python:3.11-estargz"
)

# Metrics endpoint
METRICS_URL="http://localhost:9100/metrics"

# Helper functions
log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

get_metric() {
    local metric_name=$1
    curl -s $METRICS_URL | grep "^${metric_name}" | head -1 | awk '{print $2}'
}

# Test: Cold start time
test_cold_start() {
    log_info "Testing cold start time..."
    
    local image="${TEST_IMAGES[0]}"
    
    # Clean cache
    log_info "Cleaning cache..."
    sudo rm -rf /var/lib/onyx/cache/*
    
    # Measure traditional pull
    log_info "Testing traditional pull..."
    sudo ctr images rm $image 2>/dev/null || true
    local start=$(date +%s%N)
    sudo ctr images pull $image
    local end=$(date +%s%N)
    local traditional_time=$(( (end - start) / 1000000 ))
    
    # Clean and measure lazy pull
    log_info "Testing lazy pull with onyx..."
    sudo rm -rf /var/lib/onyx/cache/*
    sudo ctr images rm $image 2>/dev/null || true
    
    start=$(date +%s%N)
    sudo ctr images pull --snapshotter onyx $image
    end=$(date +%s%N)
    local lazy_time=$(( (end - start) / 1000000 ))
    
    # Calculate improvement
    local improvement=$(( (traditional_time - lazy_time) * 100 / traditional_time ))
    
    echo "================================"
    echo "Cold Start Test Results:"
    echo "Traditional: ${traditional_time}ms"
    echo "Lazy (Onyx): ${lazy_time}ms"
    echo "Improvement: ${improvement}%"
    echo "================================"
    
    if [ $improvement -gt 50 ]; then
        log_info "✓ Cold start test PASSED (>50% improvement)"
        return 0
    else
        log_warn "✗ Cold start test WARNING (<50% improvement)"
        return 1
    fi
}

# Test: Deduplication ratio
test_dedup_ratio() {
    log_info "Testing deduplication ratio..."
    
    # Pull multiple similar images
    for image in "${TEST_IMAGES[@]}"; do
        log_info "Pulling $image..."
        sudo ctr images pull --snapshotter onyx $image 2>/dev/null || true
    done
    
    sleep 5  # Wait for metrics update
    
    # Get dedup metrics
    local dedup_ratio=$(get_metric "onyx_dedup_ratio")
    local saved_bytes=$(get_metric "onyx_dedup_saved_bytes_total")
    
    echo "================================"
    echo "Deduplication Test Results:"
    echo "Dedup Ratio: ${dedup_ratio}"
    echo "Saved Bytes: ${saved_bytes}"
    echo "================================"
    
    # Check if ratio is above threshold
    local threshold=40
    local ratio_percent=$(echo "$dedup_ratio * 100" | bc -l | cut -d. -f1)
    
    if [ "$ratio_percent" -gt "$threshold" ]; then
        log_info "✓ Dedup ratio test PASSED (>${threshold}%)"
        return 0
    else
        log_warn "✗ Dedup ratio test WARNING (<${threshold}%)"
        return 1
    fi
}

# Test: Cache hit rate
test_cache_hit() {
    log_info "Testing cache hit rate..."
    
    local image="${TEST_IMAGES[0]}"
    
    # First pull (cold)
    log_info "Cold pull..."
    sudo ctr images rm $image 2>/dev/null || true
    sudo ctr images pull --snapshotter onyx $image
    
    # Get initial metrics
    local hits_before=$(get_metric "onyx_cache_hits_total")
    local misses_before=$(get_metric "onyx_cache_misses_total")
    
    # Second pull (warm)
    log_info "Warm pull..."
    sudo ctr images rm $image 2>/dev/null || true
    sudo ctr images pull --snapshotter onyx $image
    
    sleep 2
    
    # Get final metrics
    local hits_after=$(get_metric "onyx_cache_hits_total")
    local misses_after=$(get_metric "onyx_cache_misses_total")
    
    # Calculate hit rate
    local total_hits=$(( hits_after - hits_before ))
    local total_misses=$(( misses_after - misses_before ))
    local total=$(( total_hits + total_misses ))
    
    if [ $total -eq 0 ]; then
        log_error "No cache activity detected"
        return 1
    fi
    
    local hit_rate=$(( total_hits * 100 / total ))
    
    echo "================================"
    echo "Cache Hit Rate Test Results:"
    echo "Total Requests: $total"
    echo "Cache Hits: $total_hits"
    echo "Cache Misses: $total_misses"
    echo "Hit Rate: ${hit_rate}%"
    echo "================================"
    
    if [ $hit_rate -gt 70 ]; then
        log_info "✓ Cache hit rate test PASSED (>70%)"
        return 0
    else
        log_warn "✗ Cache hit rate test WARNING (<70%)"
        return 1
    fi
}

# Test: Prefetch effectiveness
test_prefetch() {
    log_info "Testing prefetch effectiveness..."
    
    local image="${TEST_IMAGES[0]}"
    
    # Disable prefetch
    log_info "Testing without prefetch..."
    sudo sed -i 's/enabled = true/enabled = false/' /etc/onyx/config.toml
    sudo systemctl restart onyx-snapshotter
    sleep 5
    
    sudo ctr images rm $image 2>/dev/null || true
    local start=$(date +%s%N)
    sudo ctr run --rm --snapshotter onyx $image test echo "test"
    local end=$(date +%s%N)
    local time_without=$(( (end - start) / 1000000 ))
    
    # Enable prefetch
    log_info "Testing with prefetch..."
    sudo sed -i 's/enabled = false/enabled = true/' /etc/onyx/config.toml
    sudo systemctl restart onyx-snapshotter
    sleep 5
    
    sudo ctr images rm $image 2>/dev/null || true
    start=$(date +%s%N)
    sudo ctr run --rm --snapshotter onyx $image test echo "test"
    end=$(date +%s%N)
    local time_with=$(( (end - start) / 1000000 ))
    
    local improvement=$(( (time_without - time_with) * 100 / time_without ))
    
    echo "================================"
    echo "Prefetch Test Results:"
    echo "Without Prefetch: ${time_without}ms"
    echo "With Prefetch: ${time_with}ms"
    echo "Improvement: ${improvement}%"
    echo "================================"
    
    if [ $improvement -gt 10 ]; then
        log_info "✓ Prefetch test PASSED (>10% improvement)"
        return 0
    else
        log_warn "✗ Prefetch test WARNING (<10% improvement)"
        return 1
    fi
}

# Test: Memory usage
test_memory() {
    log_info "Testing memory usage..."
    
    # Get baseline memory
    local pid=$(pgrep onyx-snapshotter)
    local mem_before=$(ps -o rss= -p $pid)
    
    # Load test with multiple images
    for image in "${TEST_IMAGES[@]}"; do
        sudo ctr images pull --snapshotter onyx $image 2>/dev/null || true
    done
    
    sleep 5
    
    # Get final memory
    local mem_after=$(ps -o rss= -p $pid)
    local mem_increase=$(( mem_after - mem_before ))
    local mem_mb=$(( mem_after / 1024 ))
    
    echo "================================"
    echo "Memory Usage Test Results:"
    echo "Current Memory: ${mem_mb}MB"
    echo "Memory Increase: $(( mem_increase / 1024 ))MB"
    echo "================================"
    
    if [ $mem_mb -lt 500 ]; then
        log_info "✓ Memory test PASSED (<500MB)"
        return 0
    else
        log_warn "✗ Memory test WARNING (>500MB)"
        return 1
    fi
}

# Test: Concurrent operations
test_concurrent() {
    log_info "Testing concurrent operations..."
    
    local image="${TEST_IMAGES[0]}"
    local num_concurrent=5
    
    log_info "Starting $num_concurrent concurrent pulls..."
    
    local start=$(date +%s)
    for i in $(seq 1 $num_concurrent); do
        (
            sudo ctr images pull --snapshotter onyx $image 2>&1 | \
            grep -v "already exists" > /tmp/onyx-test-$i.log
        ) &
    done
    
    wait
    local end=$(date +%s)
    local duration=$(( end - start ))
    
    # Check for errors
    local errors=0
    for i in $(seq 1 $num_concurrent); do
        if grep -q "error\|failed" /tmp/onyx-test-$i.log; then
            errors=$(( errors + 1 ))
        fi
        rm -f /tmp/onyx-test-$i.log
    done
    
    echo "================================"
    echo "Concurrent Test Results:"
    echo "Concurrent Operations: $num_concurrent"
    echo "Duration: ${duration}s"
    echo "Errors: $errors"
    echo "================================"
    
    if [ $errors -eq 0 ]; then
        log_info "✓ Concurrent test PASSED (no errors)"
        return 0
    else
        log_error "✗ Concurrent test FAILED ($errors errors)"
        return 1
    fi
}

# Main
case "${1:-all}" in
    cold-start)
        test_cold_start
        ;;
    dedup-ratio)
        test_dedup_ratio
        ;;
    cache-hit)
        test_cache_hit
        ;;
    prefetch)
        test_prefetch
        ;;
    memory)
        test_memory
        ;;
    concurrent)
        test_concurrent
        ;;
    all)
        log_info "Running all performance tests..."
        passed=0
        failed=0
        
        test_cold_start && passed=$((passed+1)) || failed=$((failed+1))
        test_dedup_ratio && passed=$((passed+1)) || failed=$((failed+1))
        test_cache_hit && passed=$((passed+1)) || failed=$((failed+1))
        test_memory && passed=$((passed+1)) || failed=$((failed+1))
        test_concurrent && passed=$((passed+1)) || failed=$((failed+1))
        
        echo ""
        echo "================================"
        echo "Test Summary:"
        echo "Passed: $passed"
        echo "Failed: $failed"
        echo "================================"
        
        [ $failed -eq 0 ] && exit 0 || exit 1
        ;;
    *)
        echo "Usage: $0 {cold-start|dedup-ratio|cache-hit|prefetch|memory|concurrent|all}"
        exit 1
        ;;
esac